package com.gome.ai.aiparter.ws;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.gome.ai.aiparter.ws.em.ActionEnum;
import com.gome.ai.aiparter.ws.em.AiStateEnum;
import com.gome.ai.aiparter.ws.enity.Answer;
import com.gome.ai.aiparter.ws.enity.WSResponse;
import lombok.extern.slf4j.Slf4j;

import java.util.concurrent.LinkedBlockingDeque;
import java.util.function.Consumer;

/**
 * AI 客户端可接收输入
 */
@Slf4j
public class AIConnection implements WSConnection {
    //空闲AI
    public static final LinkedBlockingDeque<AIConnection> idleAIs = new LinkedBlockingDeque<>();

    protected final WebSocketServer server;

    protected Consumer<String> messageConsumer;

    public AIConnection(WebSocketServer server) {
        this.server = server;
    }

    @Override
    public void onMessage(String message) {
        WSResponse res = JSON.parseObject(message, WSResponse.class);
        if (res.getAct() == ActionEnum.ANSWER.getCode()) {
            if(messageConsumer != null){
                messageConsumer.accept(message);
            }

            JSONObject data = (JSONObject) res.getData();
            Answer answer = data.toJavaObject(Answer.class);
            onAnswer(answer,res,message);
        }else {
            log.warn("AI未知消息:{}",message);
        }
    }

    protected void onAnswer(Answer answer,WSResponse response,String message){
        if (answer.getState() == AiStateEnum.READY.getCode()) {
            log.info("AI准备就绪 token = {}",server.getToken());
            messageConsumer = null;
            idleAIs.add(this);
            log.debug("空闲AI数量 size = {}",idleAIs.size());
        }else if(answer.getState() == AiStateEnum.STOP.getCode()){
            log.info("AI停止 token = {}",server.getToken());
            idleAIs.remove(this);
        }
    }

    @Override
    public void onClose() {
        boolean remove = idleAIs.remove(this);
        if(remove){
            log.info("AI连接关闭,移除空闲队列 token = {}",server.getToken());
        }else {
            log.warn("AI连接关闭,AI没有在空闲队列中");
        }

        log.debug("空闲AI数量 size = {}",idleAIs.size());
    }

    /**
     * 消费AI的答案
     * @param messageConsumer
     */
    public void answer(Consumer<String> messageConsumer) {
        this.messageConsumer = messageConsumer;
    }

    /**
     * 向AI提问题
     * @param message
     */
    public void ask(String message) {
        server.send(message);
    }
}
